package cn.ddiancan.auth.service.generator;

import cn.ddiancan.auth.constant.RequestSourceEnum;
import cn.ddiancan.auth.service.userdetails.XddcloudUserDetailsImpl;
import cn.ddiancan.auth.utils.JwtUtils;
import org.springframework.lang.Nullable;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.settings.OAuth2TokenFormat;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;

public class XddJwtGenerator implements OAuth2TokenGenerator<Jwt> {

    private static final String REQUEST_SOURCE = "requestSource";

    private final JwtEncoder jwtEncoder;

    private final OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer;

    private final XddcloudUserDetailsImpl userDetailsService;

    public XddJwtGenerator(JwtEncoder jwtEncoder, OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer,
        XddcloudUserDetailsImpl userDetailsService) {
        Assert.notNull(jwtCustomizer, "jwtCustomizer cannot be null");
        Assert.notNull(jwtEncoder, "jwtEncoder cannot be null");
        Assert.notNull(userDetailsService, "userDetailsService cannot be null");
        this.userDetailsService = userDetailsService;
        this.jwtCustomizer = jwtCustomizer;
        this.jwtEncoder = jwtEncoder;
    }

    @Nullable
    @Override
    public Jwt generate(OAuth2TokenContext context) {
        if (context.getTokenType() == null || (!OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType()))) {
            return null;
        }
        if (OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType()) && !OAuth2TokenFormat.SELF_CONTAINED.equals(
            context.getRegisteredClient().getTokenSettings().getAccessTokenFormat())) {
            return null;
        }
        String issuer = null;
        if (context.getAuthorizationServerContext() != null) {
            issuer = context.getAuthorizationServerContext().getIssuer();
        }
        RegisteredClient registeredClient = context.getRegisteredClient();
        JwtClaimsSet.Builder claimsBuilder =
            JwtUtils.accessTokenClaims(registeredClient, issuer, context.getPrincipal().getName(),
                context.getAuthorizedScopes());
        JwsHeader.Builder jwsHeaderBuilder = JwtUtils.headers();
        JwtEncodingContext jwtEncodingContext =
            JwtEncodingContext.with(jwsHeaderBuilder, claimsBuilder).registeredClient(registeredClient)
                .principal(context.getPrincipal()).authorization(context.getAuthorization())
                .authorizedScopes(context.getAuthorizedScopes()).tokenType(OAuth2TokenType.ACCESS_TOKEN)
                .authorizationGrantType(AuthorizationGrantType.PASSWORD)
                .authorizationGrant(context.getAuthorizationGrant()).build();
        userDetailsService.setRequestSource(RequestSourceEnum.valueOf(context.get(REQUEST_SOURCE)));
        userDetailsService.loadUserByUsername(context.getPrincipal().getName());
        if (this.jwtCustomizer != null) {
            jwtCustomizer.customize(jwtEncodingContext);
        }
        JwsHeader jwsHeader = jwsHeaderBuilder.build();
        JwtClaimsSet claims = claimsBuilder.build();
        return this.jwtEncoder.encode(JwtEncoderParameters.from(jwsHeader, claims));
    }
}
